import json
import os
import os.path as osp
import time

import cv2
import numpy as np
from tqdm import tqdm


def get_img_list(path):
    img_list = []
    for root, dirs, files in os.walk(path):
        for f in files:
            if f.endswith('.jpg') and f.find('_mask') == -1:
                img_list.append(osp.join(root, f))

    assert len(img_list) > 0, "没有图片！"
    return img_list

    
def get_heatmap(img, contours):
    # tic = time.time()
    res = np.zeros_like(img[:, :, 0], dtype=np.float32)

    for contour in contours:
        # contour[i] ---> (x, y)
        mask = np.zeros_like(img[:,:,0], dtype=np.float32)
        # 面
        # cv2.fillPoly(mask, [contour], 1)
        cv2.fillPoly(mask, [np.around(contour).astype(np.int64)], 1)

        # # 线
        # cv2.polylines(mask, [contour], False, 1, thickness=25)

        where_condition = np.where(mask==1)
        x0 = where_condition[1]
        y0 = where_condition[0]

        ds = []
        pnum = len(contour)
        for i in range(pnum):
            p1 = contour[i]
            p2 = contour[(i + 1) % pnum]

            x1 = np.ones_like(x0, dtype=np.float32) * p1[0]
            y1 = np.ones_like(y0, dtype=np.float32) * p1[1]
            x2 = np.ones_like(x0, dtype=np.float32) * p2[0]
            y2 = np.ones_like(y0, dtype=np.float32) * p2[1]

            a = y2 - y1
            b = x1 - x2
            c = x2 * y1 - x1 * y2
            molecule = np.abs(a * x0 + b * y0 + c)

            # 根据各情况计算垂足点
            xfoot = np.zeros_like(x0, dtype=np.float32)
            yfoot = np.zeros_like(y0, dtype=np.float32)
            other_condition = np.ones_like(x0, dtype=np.int64)

            # 情况一：p1和p2两点重合，故垂足点即为p1(或p2)
            condition = a ** 2 + b ** 2 < 1e-10
            xfoot[condition] = x1[condition]
            yfoot[condition] = y1[condition]
            other_condition[condition] = 0

            # 情况二： p0点在线p1-p2上，故垂足点即为p0
            condition = molecule < 1e-10
            xfoot[condition] = x0[condition]
            yfoot[condition] = y0[condition]
            other_condition[condition] = 0

            # 其余情况：
            other_condition = other_condition > 0
            xfoot[other_condition] = (b[other_condition]**2 * x0[other_condition] - a[other_condition] * b[other_condition] * y0[other_condition] - \
                                          a[other_condition] * c[other_condition])/(a[other_condition]**2 + b[other_condition]**2)
            yfoot[other_condition] = (-a[other_condition] * b[other_condition] * x0[other_condition] + a[other_condition]**2 * y0[other_condition] - \
                                          b[other_condition] * c[other_condition])/(a[other_condition]**2 + b[other_condition]**2)

            # 计算点到线段的距离
            distance = np.zeros_like(x0, dtype=np.float32)

            tmp_x = np.c_[x1, x2]
            tmp_y = np.c_[y1, y2]
            condition1 = np.min(tmp_x, -1) <= xfoot
            condition2 = xfoot <= np.max(tmp_x, -1)
            condition3 = np.min(tmp_y, -1) <= yfoot
            condition4 = yfoot <= np.max(tmp_y, -1)

            # np.min(tmp_x, -1) <= xfoot <= np.max(tmp_x, -1) and np.min(tmp_y, -1) <= yfoot <= np.max(tmp_y, -1)
            # 垂足点在线段p1-p2上：
            condition = condition1 & condition2 & condition3 & condition4
            if (~(np.sqrt(a[condition] ** 2 + b[condition] ** 2) > 0)).any():
                print("error !!!!!!!!!")
                exit()
            distance[condition] = molecule[condition] / np.sqrt(a[condition] ** 2 + b[condition] ** 2)
            # 垂足点在线段p1-p2外，则取p0分别到p1和p2的最小距离
            condition = ~condition
            distance[condition] = np.min(np.c_[np.sqrt((x0[condition] - x1[condition]) ** 2 + (y0[condition] - y1[condition]) ** 2),  \
                                               np.sqrt((x0[condition] - x2[condition]) ** 2 + (y0[condition] - y2[condition]) ** 2)], -1)
            ds.append(distance)

        distance = np.stack(ds, -1)
        min_distance = np.min(distance, -1)
        maxmin_distance = np.max(min_distance)

        # 面
        res[where_condition] = min_distance / maxmin_distance

        # # 线
        # res[where_condition] = 1 -  min_distance / maxmin_distance

    # toc = time.time()
    # print(toc - tic)
    return res

def label_process(image, label_path):
    with open(label_path, 'r') as f:
        content = json.loads(f.read())


    contours = []
    for shape in content['shapes']:
        points = shape['points']
        contours.append(np.array(points, dtype=np.int64))

    heat_img = get_heatmap(image, contours)
    
    return heat_img


if __name__ == '__main__':
    rootpth = r'data\shushu\train'
    save_path = r'data\shushu\train_heatmap'
    os.makedirs(save_path, exist_ok=True)

    img_list = get_img_list(rootpth)    #[1:2]
    total_data = []
    for im_path in img_list:
        label_path = im_path.replace('.jpg', '.json')
        assert os.path.exists(label_path), "path is not exist"
        total_data.append([im_path, label_path])

    for i in tqdm(range(len(total_data))):
        img_path, anno_path = total_data[i]
        print("img path: {}".format(img_path))
        # dir_path = osp.dirname(img_path)
        base_name = osp.basename(img_path)
        mask_name = base_name.replace('.jpg', '_mask.png')
        mask_path = osp.join(save_path, mask_name)

        img = cv2.imread(img_path)
        # print(img.shape)

        with open(anno_path, 'r') as f:
            content = json.loads(f.read())


        contours = []
        for shape in content['shapes']:
            points = shape['points']
            points = np.array(points, dtype=np.float32)
            contours.append(points)

        heat_img = get_heatmap(img, contours)
        heat_img = np.array(heat_img*255, np.uint8)


        # cv2.imshow('img', img)

        cv2.imwrite(mask_path, heat_img)
        # cv2.imshow('heat', heat_img)
        # cv2.waitKey()
